import os
import sys
import argparse
from Crypto.Cipher import AES
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Random import get_random_bytes

class Encryption:
    @staticmethod
    def generate_random_salt():
        """生成32字节的随机盐值"""
        return get_random_bytes(32)

    @staticmethod
    def file_encrypt(input_file, password):
        """加密文件"""
        try:
            # 生成盐值
            salt = Encryption.generate_random_salt()
            
            # 创建输出文件
            output_file = input_file + ".Abantes"
            
            # 准备AES加密参数
            key_size = 256
            block_size = 128
            
            # 使用PBKDF2生成密钥和IV
            key_iv = PBKDF2(password.encode('utf-8'), salt, dkLen=(key_size + block_size) // 8, count=50000)
            key = key_iv[:key_size // 8]
            iv = key_iv[key_size // 8:key_size // 8 + block_size // 8]
            
            # 创建AES加密器
            cipher = AES.new(key, AES.MODE_CFB, iv=iv, segment_size=128)
            
            with open(input_file, 'rb') as file_in:
                with open(output_file, 'wb') as file_out:
                    # 写入盐值
                    file_out.write(salt)
                    
                    # 加密并写入数据
                    while True:
                        chunk = file_in.read(1048576)  # 1MB chunks
                        if not chunk:
                            break
                        encrypted_chunk = cipher.encrypt(chunk)
                        file_out.write(encrypted_chunk)
                        
            print(f"文件加密成功: {output_file}")
            return True
            
        except Exception as e:
            print(f"加密错误: {e}")
            return False

    @staticmethod
    def file_decrypt(input_file, output_file, password):
        """解密文件"""
        try:
            with open(input_file, 'rb') as file_in:
                # 读取盐值
                salt = file_in.read(32)
                
                # 准备AES解密参数
                key_size = 256
                block_size = 128
                
                # 使用PBKDF2生成密钥和IV
                key_iv = PBKDF2(password.encode('utf-8'), salt, dkLen=(key_size + block_size) // 8, count=50000)
                key = key_iv[:key_size // 8]
                iv = key_iv[key_size // 8:key_size // 8 + block_size // 8]
                
                # 创建AES解密器
                cipher = AES.new(key, AES.MODE_CFB, iv=iv, segment_size=128)
                
                with open(output_file, 'wb') as file_out:
                    # 解密并写入数据
                    while True:
                        chunk = file_in.read(1048576)  # 1MB chunks
                        if not chunk:
                            break
                        decrypted_chunk = cipher.decrypt(chunk)
                        file_out.write(decrypted_chunk)
                        
            print(f"文件解密成功: {output_file}")
            
            #Debug： 验证解密后的文件
            #if output_file.lower().endswith('.mp3'):
            #    with open(output_file, 'rb') as f:
            #        header = f.read(3)
            #        if header == b'ID3':
            #            print("MP3文件头验证成功")
            #        else:
            #            print(f"警告: MP3文件头异常，前3字节: {header.hex()}")
            #
            #return True
            
        except Exception as e:
            print(f"解密错误: {e}")
            return False

def main():
    parser = argparse.ArgumentParser(description='文件加密解密工具')
    parser.add_argument('-encrypt', metavar='文件路径', help='加密文件')
    parser.add_argument('-decrypt', metavar='文件路径', help='解密文件')
    parser.add_argument('key', nargs='?', default='WR8h2GIbf9FGz6VVlSzJ', help='加密密钥（可选，默认为WR8h2GIbf9FGz6VVlSzJ）')
    
    args = parser.parse_args()
    
    # 安装依赖检查
    try:
        from Crypto.Cipher import AES
    except ImportError:
        print("请安装pycryptodome库: pip install pycryptodome")
        return
    
    if args.encrypt:
        if not os.path.exists(args.encrypt):
            print(f"文件不存在: {args.encrypt}")
            return
            
        Encryption.file_encrypt(args.encrypt, args.key)
        
    elif args.decrypt:
        if not os.path.exists(args.decrypt):
            print(f"文件不存在: {args.decrypt}")
            return
            
        # 生成解密后的文件名
        if args.decrypt.endswith('.Abantes'):
            output_file = args.decrypt[:-8]  # 移除.Abantes后缀
        else:
            output_file = args.decrypt + '.decrypted'
            
        # 检查输出文件是否已存在
        if os.path.exists(output_file):
            response = input(f"文件 {output_file} 已存在，是否覆盖？(y/n): ")
            if response.lower() != 'y':
                print("操作取消")
                return
        
        Encryption.file_decrypt(args.decrypt, output_file, args.key)
        
    else:
        print("请指定 -encrypt 或 -decrypt 参数")
        parser.print_help()

if __name__ == "__main__":
    main()